import pandas as pd
import argparse
import os
from tqdm import tqdm

def parquet_to_text(input_file, output_file, format='csv', sep=',', encoding='utf-8'):
    """
    将Parquet文件转换为文本文件
    
    参数:
        input_file (str): 输入的Parquet文件路径
        output_file (str): 输出的文本文件路径
        format (str): 输出格式，支持'csv'、'txt'等
        sep (str): 分隔符，默认为逗号
        encoding (str): 输出文件编码，默认为utf-8
    """
    print(f"读取Parquet文件: {input_file}")
    # 读取Parquet文件
    try:
        df = pd.read_parquet(input_file)
    except Exception as e:
        print(f"读取Parquet文件出错: {e}")
        return False
    
    print(f"Parquet文件包含 {len(df)} 行和 {len(df.columns)} 列")
    
    # 根据指定格式保存为文本文件
    try:
        if format.lower() == 'csv':
            df.to_csv(output_file, sep=sep, index=False, encoding=encoding)
        elif format.lower() == 'txt':
            df.to_csv(output_file, sep=sep, index=False, header=False, encoding=encoding)
        else:
            # 默认保存为CSV
            df.to_csv(output_file, sep=sep, index=False, encoding=encoding)
        
        print(f"文件已成功转换并保存至: {output_file}")
        return True
    except Exception as e:
        print(f"保存文件时出错: {e}")
        return False

def batch_convert(input_dir, output_dir, format='csv', sep=',', encoding='utf-8'):
    """
    批量转换目录中的所有Parquet文件
    
    参数:
        input_dir (str): 输入目录，包含Parquet文件
        output_dir (str): 输出目录
        format (str): 输出格式
        sep (str): 分隔符
        encoding (str): 输出文件编码
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取所有parquet文件
    parquet_files = [f for f in os.listdir(input_dir) if f.endswith('.parquet')]
    
    if not parquet_files:
        print(f"目录 {input_dir} 中未找到Parquet文件")
        return
    
    print(f"找到 {len(parquet_files)} 个Parquet文件，开始转换...")
    
    # 使用tqdm显示进度
    for file in tqdm(parquet_files):
        input_path = os.path.join(input_dir, file)
        output_path = os.path.join(output_dir, f"{os.path.splitext(file)[0]}.{format}")
        parquet_to_text(input_path, output_path, format, sep, encoding)

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='将Parquet文件转换为文本文件')
    
    # 添加命令行参数
    parser.add_argument('--input', '-i', required=True, help='输入的Parquet文件或目录')
    parser.add_argument('--output', '-o',  default='dataset/output.java', help='输出的文本文件或目录')
    parser.add_argument('--format', '-f', default='txt', help='输出格式 (csv, txt)')
    parser.add_argument('--separator', '-s', default=',', help='字段分隔符')
    parser.add_argument('--encoding', '-e', default='utf-8', help='输出文件编码')
    parser.add_argument('--batch', '-b', action='store_true', help='批量处理目录中的所有Parquet文件')
    
    args = parser.parse_args()
    
    if args.batch:
        batch_convert(args.input, args.output, args.format, args.separator, args.encoding)
    else:
        parquet_to_text(args.input, args.output, args.format, args.separator, args.encoding)

if __name__ == '__main__':
    main() 